import json
from transformers import AutoTokenizer
from collections import Counter
import numpy as np

# 载入 Qwen2.5-32B 的分词器
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-32B")

input_lengths = []
output_lengths = []

with open("600/train.jsonl", "r", encoding="utf-8") as f:
    for line in f:
        data = json.loads(line)
        instr = data.get("instruction", "")
        inp   = data.get("input", "")
        # 拼接 instruction 和 input
        text  = instr + (" " + inp if inp else "")
        out   = data.get("output", "")

        # 分词并计数（不含 special tokens）
        input_ids  = tokenizer.encode(text, add_special_tokens=False)
        output_ids = tokenizer.encode(out,  add_special_tokens=False)
        input_lengths.append(len(input_ids))
        output_lengths.append(len(output_ids))

# 计算常用统计量
def calc_stats(lengths):
    return {
        "count":  len(lengths),
        "min":    int(np.min(lengths)),
        "max":    int(np.max(lengths)),
        "mean":   float(np.mean(lengths)),
        "median": float(np.median(lengths)),
        "p90":    float(np.percentile(lengths, 90)),
        "p95":    float(np.percentile(lengths, 95)),
    }

input_stats  = calc_stats(input_lengths)
output_stats = calc_stats(output_lengths)

print("=== Input Token Length Stats ===")
for k, v in input_stats.items():
    print(f"{k:>6} :", v)

print("\n=== Output Token Length Stats ===")
for k, v in output_stats.items():
    print(f"{k:>6} :", v)

# 如果你还想可视化直方图，可以取消下面代码的注释：
import matplotlib.pyplot as plt
plt.hist(input_lengths, bins=50)
plt.title("Input Length Distribution")
plt.xlabel("Token Count")
plt.ylabel("Frequency")

# 保存到当前工作目录
plt.savefig("input_length_hist.png", dpi=150, bbox_inches="tight")
plt.close()

# 同理，输出长度：
plt.hist(output_lengths, bins=50)
plt.title("Output Length Distribution")
plt.xlabel("Token Count")
plt.ylabel("Frequency")
plt.savefig("output_length_hist.png", dpi=150, bbox_inches="tight")
plt.close()
